package com.iscas.loadstorm.isloadstorm.util;

import com.iscas.loadstorm.isloadstorm.collector.CustomResultCollector;
import com.iscas.loadstorm.isloadstorm.domain.TestRun;
import org.apache.jmeter.engine.StandardJMeterEngine;

import org.apache.jmeter.save.SaveService;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.collections.HashTree;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class JmxRunner {
    /**
     * 运行 JMeter 测试
     *
     * @param loadscript Loadscript 对象，包含脚本 URL
     * @return 运行状态消息
     */

    @Value("${jmeter.home}")
    private String jmeterHome;

    // 存储测试任务状态
    private static final ConcurrentHashMap<String, TestRun> testRuns = new ConcurrentHashMap<>();

    // 初始化 JMeter 环境
    private void initJMeter() {
        System.out.println(jmeterHome);
        JMeterUtils.setJMeterHome(jmeterHome);
        JMeterUtils.loadJMeterProperties(jmeterHome + "/bin/jmeter.properties");
        JMeterUtils.initLocale();
    }

    /**
     * 启动 JMeter 压测任务
     *
     * @param loadScriptUrl JMX 文件 URL
     * @return 测试任务的唯一 ID
     */
    public String runJMeterTest(String loadScriptUrl) {
        String testId = UUID.randomUUID().toString();

        try {
            initJMeter();
            String jmxFilePath = downloadJmxFile(loadScriptUrl);
            if (jmxFilePath == null) {
                throw new RuntimeException("Failed to download JMX file.");
            }

            File jmxFile = new File(jmxFilePath);
            HashTree testPlanTree = SaveService.loadTree(jmxFile);

            StandardJMeterEngine jmeterEngine = new StandardJMeterEngine();

            // 先创建 TestRun 实例,暂时传入 null 作为 resultCollector
            TestRun testRun = new TestRun(testId, jmeterEngine, null);
            testRuns.put(testId, testRun);

            // 创建 CustomResultCollector 并传入 testRun
            CustomResultCollector resultCollector = new CustomResultCollector(testRun);

            // 更新 TestRun 的 resultCollector
            testRun.setResultCollector(resultCollector);

            // 将结果收集器添加到测试计划中
            testPlanTree.add(testPlanTree.getArray()[0], resultCollector);

            new Thread(() -> {
                try {
                    jmeterEngine.configure(testPlanTree);
                    jmeterEngine.run();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }).start();

            return testId;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("JMeter test failed to start: " + e.getMessage());
        }
    }
    /**
     * 下载 JMX 文件到本地临时路径
     *
     * @param fileUrl JMX 文件的 URL
     * @return 本地 JMX 文件路径
     * @throws IOException 如果下载失败
     */
    private static String downloadJmxFile(String fileUrl) throws IOException {
        String tempFilePath = "/tmp/temp_jmx_file.jmx";
        HttpURLConnection connection = null;
        try (InputStream inputStream = new URL(fileUrl).openStream();
             FileOutputStream outputStream = new FileOutputStream(tempFilePath)) {

            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            return tempFilePath;
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        } finally {
            if (connection != null) {
                connection.disconnect();
            }
        }
    }

    public static String generateResultFileName(String basePath) {
        String timestamp = new SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date());
        return basePath + "/test_results_" + timestamp + ".csv";
    }


    /**
     * 获取指定测试ID的TestRun
     *
     * @param testId 测试ID
     * @return TestRun对象，如果不存在返回null
     */
    public static TestRun getTestRun(String testId) {
        return testRuns.get(testId);
    }

    /**
     * 移除指定测试ID的TestRun
     *
     * @param testId 测试ID
     */
    public static void removeTestRun(String testId) {
        testRuns.remove(testId);
    }

    /**
     * 停止指定的测试
     *
     * @param testId 测试ID
     */
    public static void stopTest(String testId) {
        TestRun testRun = testRuns.get(testId);
        if (testRun != null && testRun.getEngine() != null) {
            testRun.getEngine().stopTest(true);
            removeTestRun(testId);
        }
    }
}
